package trie

type TrieNode struct{
	Word rune // 当前节点储存的字符，byte只能表示英文字符，rune可以表示任意字符
	Children map[rune]*TrieNode // 子节点
	Term string // 如果word是一个Trem的结尾，则保存这个Trem
}

//trie树
type TrieTree struct{
	root *TrieNode
}

// 向trie树中插入一个词
func(tree *TrieTree) AddTerm(term string){
	if len(term) <= 1{
		return 
	}
	word := []rune(term)

	if tree.root == nil{
		tree.root = new(TrieNode) // 初始化根节点
	}

	tree.root.add(word,term,0)
}

// add 把word[beginIndex:]插入到树中
func(node *TrieNode)add(word []rune,term string,beginIndex int){
	if beginIndex >= len(word){ //word 已经遍历完
		node.Term = term
		return 
	}

	wordRune := word[beginIndex] // 当前要处理的字符,放入到子节点中

	if node.Children == nil{
		node.Children = make(map[rune]*TrieNode)
	}

	if child,exists := node.Children[wordRune]; exists{
		newNode := &TrieNode{Word: wordRune}
		node.Children[wordRune] = newNode
		child.add(word,term,beginIndex+1)
	}else{
		child.add(word,term,beginIndex+1)
	}
}

// 检索满足要求的字符串是否在树中
func(tree *TrieTree) Retireve(prefix string)  []string{
	if tree.root == nil || len(prefix) <= 0{
		return nil
	}
	words := []rune(prefix)
	firstWord := words[0]
	if child,exists := tree.root.Children[firstWord]; exists{
		end := child.walk(words,0)
		if end == nil{
			return nil
		}else{
			trems := make([]string,0,100)
			end.traverseTerms(&trems)
			return trems
		}
	}
	return nil
}

// work上一个节点存储的字符，如果当前节点是叶子节点，则返回叶子节点的Term，否则返回nil
func(node *TrieNode) walk(words []rune,beginIndex int) *TrieNode{
	if beginIndex == len(words) - 1{
		return node
	}
	beginIndex += 1
	word := words[beginIndex]
	if child,exists := node.Children[word]; exists{
		return child.walk(words,beginIndex)
	}else{
		return nil
	}
}

func (node *TrieNode) traverseTerms(term *[]string){
	if len(node.Term) > 0{
		*term = append(*term,node.Term)
	}
	for _,child := range node.Children{
		child.traverseTerms(term)
	}
}